import numpy as np
import matplotlib.pyplot as plt
import random
import math

from Rosenbroke import *

def trustRegionMethod(x,func,dfunc,ddfunc,imax):
    """
    信任域方法，dogleg策略:
    param:{
        x:初值
        func：函数
        dfunc：函数一阶偏导
        ddfunc：函数二阶偏导
        imax:迭代次数
    }
    """
    x0 = x
    deltaMax = 10
    delta = 0.1
    eta = 0
    W=np.zeros((x0.shape[0],imax))
    W[:,0] = x0

    for i in range(imax):
        fk = func(x0)
        dfk = dfunc(x0)
        Bk = ddfunc(x0)
        sk = s(dfk,Bk,delta)

        rhok = rho_k(func,dfunc,ddfunc,m,x,sk)
        
        if rhok < 1/4 :
            delta = 1/4*delta
        elif rhok>3/4 and (abs(math.sqrt(np.dot(sk,sk))-delta) < 10e-10):
            delta = min(2*delta,deltaMax)
        
        if rhok > eta:
            x0 = x0 + sk
        W[:,i] = x0

    return W

def m(funcValue,dFuncValue,p,Bk):
    ans = funcValue + np.dot(dFuncValue,p) + 1/2 * np.dot(np.dot(p.T,Bk),p)
    return ans



def tau(pU,pB,delta):
    npB = math.sqrt(np.dot(pB.T,pB))
    npU = math.sqrt(np.dot(pU.T,pU))
    
    if npB <=delta:
        tau=2
    elif npU >= delta:
        tau = delta/npU
    else:
        pB_U = pB - pU
        tau = (-np.dot(pU.T,pB_U) + math.sqrt((np.dot(pU.T,pB_U)**2 - np.dot(pB_U.T,pB_U)*(np.dot(pU.T,pU)-delta**2))))/(np.dot(pB_U.T,pB_U))
        tau = tau+1

    
    return tau


def s(dfunValue, Hessian, delta):
    pU = - (np.dot(dfunValue,dfunValue)/(np.dot(np.dot(dfunValue.T,Hessian),dfunValue))) * dfunValue
    pB = - np.dot((np.linalg.inv(Hessian)),dfunValue)
    taU = tau(pU,pB,delta)
    
    if taU >= 0 and taU <= 1 :
        sk = taU * pU
    elif taU >= 1 and taU <= 2:
        sk = pU + (taU-1)*(pB - pU)
    return sk 


def rho_k(func,dfunc,ddfunc,m,x,sk):
    rho = (func(x)-func(x+sk))/(m(func(x),dfunc(x),np.array([0,0]),ddfunc(x))-m(func(x),dfunc(x),sk,ddfunc(x)))
    
    return rho


if __name__ == "__main__":
    X1=np.arange(-10,10+0.05,0.05)
    X2=np.arange(-10,10+0.05,0.05)
    [x1,x2]=np.meshgrid(X1,X2)
    f = (100*((x2-x1**2)**2)+(x1-1)**2)
    plt.contour(x1,x2,f,20) # 画出函数的20条轮廓线
    temp = [-7.0 , 0.5]
    x0 = np.array(temp)
    imax = 1000
    W = trustRegionMethod(x0,Rosenbroke,dRosenbroke,ddRosenbroke,imax)
    print("迭代结束点：")
    print(W[:,imax-1])
    print("迭代结果：")
    print(Rosenbroke(W[:,imax-1]))
    plt.plot(W[0,:],W[1,:],'g*',W[0,:],W[1,:]) # 画出迭代点收敛的轨迹
    plt.show()